{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SAC2018.ipynb",
      "version": "0.3.2",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Bp5fR_MSBLe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.distributions import Normal\n",
        "\n",
        "import gym\n",
        "import random\n",
        "import numpy as np\n",
        "from collections import deque"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t6LsUoK3SGTR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def mini_batch_train(env, agent, max_episodes, max_steps, batch_size):\n",
        "    episode_rewards = []\n",
        "    update_step = 0\n",
        "\n",
        "    for episode in range(max_episodes):\n",
        "        state = env.reset()\n",
        "        episode_reward = 0\n",
        "\n",
        "        for step in range(max_steps):\n",
        "            action = agent.get_action(state)\n",
        "            next_state, reward, done, _ = env.step(action)\n",
        "            agent.replay_buffer.push(state, action, reward, next_state, done)\n",
        "            episode_reward += reward\n",
        "\n",
        "            if len(agent.replay_buffer) > batch_size:\n",
        "                agent.update(batch_size)\n",
        "                update_step += 1\n",
        "\n",
        "            if done or step == max_steps-1:\n",
        "                episode_rewards.append(episode_reward)\n",
        "                break\n",
        "\n",
        "            state = next_state\n",
        "        \n",
        "        print(\"Episode \" + str(episode) + \": \" + str(episode_reward))\n",
        "\n",
        "    return episode_rewards"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FZSveC3rSHBb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class BasicBuffer:\n",
        "\n",
        "  def __init__(self, max_size):\n",
        "      self.max_size = max_size\n",
        "      self.buffer = deque(maxlen=max_size)\n",
        "\n",
        "  def push(self, state, action, reward, next_state, done):\n",
        "      experience = (state, action, np.array([reward]), next_state, done)\n",
        "      self.buffer.append(experience)\n",
        "\n",
        "  def sample(self, batch_size):\n",
        "      state_batch = []\n",
        "      action_batch = []\n",
        "      reward_batch = []\n",
        "      next_state_batch = []\n",
        "      done_batch = []\n",
        "\n",
        "      batch = random.sample(self.buffer, batch_size)\n",
        "\n",
        "      for experience in batch:\n",
        "          state, action, reward, next_state, done = experience\n",
        "          state_batch.append(state)\n",
        "          action_batch.append(action)\n",
        "          reward_batch.append(reward)\n",
        "          next_state_batch.append(next_state)\n",
        "          done_batch.append(done)\n",
        "\n",
        "      return (state_batch, action_batch, reward_batch, next_state_batch, done_batch)\n",
        "\n",
        "  def __len__(self):\n",
        "      return len(self.buffer)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VjrRWx-2Sx1p",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ValueNetwork(nn.Module):\n",
        "\n",
        "    def __init__(self, input_dim, output_dim, init_w=3e-3):\n",
        "        super(ValueNetwork, self).__init__()\n",
        "        self.fc1 = nn.Linear(input_dim, 256)\n",
        "        self.fc2 = nn.Linear(256, 256)\n",
        "        self.fc3 = nn.Linear(256, output_dim)\n",
        "\n",
        "        self.fc3.weight.data.uniform_(-init_w, init_w)\n",
        "        self.fc3.bias.data.uniform_(-init_w, init_w)\n",
        "\n",
        "    def forward(self, state):\n",
        "        x = F.relu(self.fc1(state))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "class SoftQNetwork(nn.Module):\n",
        "    \n",
        "    def __init__(self, num_inputs, num_actions, hidden_size=256, init_w=3e-3):\n",
        "        super(SoftQNetwork, self).__init__()\n",
        "        self.linear1 = nn.Linear(num_inputs + num_actions, hidden_size)\n",
        "        self.linear2 = nn.Linear(hidden_size, hidden_size)\n",
        "        self.linear3 = nn.Linear(hidden_size, 1)\n",
        "\n",
        "        self.linear3.weight.data.uniform_(-init_w, init_w)\n",
        "        self.linear3.bias.data.uniform_(-init_w, init_w)\n",
        "\n",
        "    def forward(self, state, action):\n",
        "        x = torch.cat([state, action], 1)\n",
        "        x = F.relu(self.linear1(x))\n",
        "        x = F.relu(self.linear2(x))\n",
        "        x = self.linear3(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "class PolicyNetwork(nn.Module):\n",
        "    \n",
        "    def __init__(self, num_inputs, num_actions, hidden_size=256, init_w=3e-3, log_std_min=-20, log_std_max=2):\n",
        "        super(PolicyNetwork, self).__init__()\n",
        "        self.log_std_min = log_std_min\n",
        "        self.log_std_max = log_std_max\n",
        "\n",
        "        self.linear1 = nn.Linear(num_inputs, hidden_size)\n",
        "        self.linear2 = nn.Linear(hidden_size, hidden_size)\n",
        "\n",
        "        self.mean_linear = nn.Linear(hidden_size, num_actions)\n",
        "        self.mean_linear.weight.data.uniform_(-init_w, init_w)\n",
        "        self.mean_linear.bias.data.uniform_(-init_w, init_w)\n",
        "\n",
        "        self.log_std_linear = nn.Linear(hidden_size, num_actions)\n",
        "        self.log_std_linear.weight.data.uniform_(-init_w, init_w)\n",
        "        self.log_std_linear.bias.data.uniform_(-init_w, init_w)\n",
        "\n",
        "    def forward(self, state):\n",
        "        x = F.relu(self.linear1(state))\n",
        "        x = F.relu(self.linear2(x))\n",
        "\n",
        "        mean    = self.mean_linear(x)\n",
        "        log_std = self.log_std_linear(x)\n",
        "        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)\n",
        "\n",
        "        return mean, log_std\n",
        "\n",
        "    def sample(self, state, epsilon=1e-6):\n",
        "        mean, log_std = self.forward(state)\n",
        "        std = log_std.exp()\n",
        "\n",
        "        normal = Normal(mean, std)\n",
        "        z = normal.rsample()\n",
        "        action = torch.tanh(z)\n",
        "\n",
        "        log_pi = normal.log_prob(z) - torch.log(1 - action.pow(2) + epsilon)\n",
        "        log_pi = log_pi.sum(1, keepdim=True)\n",
        "\n",
        "        return action, log_pi"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vQJka_ocSLSX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class SACAgent:\n",
        "  \n",
        "    def __init__(self, env, gamma, tau, v_lr, q_lr, policy_lr, buffer_maxlen):\n",
        "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "        \n",
        "        self.env = env        \n",
        "        self.action_range = [env.action_space.low, env.action_space.high]\n",
        "\n",
        "        # hyperparameters\n",
        "        self.gamma = gamma\n",
        "        self.tau = tau       \n",
        "        self.update_step = 0\n",
        "        self.delay_step = 2\n",
        "\n",
        "        # initialize networks \n",
        "        self.value_net = ValueNetwork(env.observation_space.shape[0], 1).to(self.device)\n",
        "        self.target_value_net = ValueNetwork(env.observation_space.shape[0], 1).to(self.device)\n",
        "        self.q_net1 = SoftQNetwork(env.observation_space.shape[0], env.action_space.shape[0]).to(self.device)\n",
        "        self.q_net2 = SoftQNetwork(env.observation_space.shape[0], env.action_space.shape[0]).to(self.device)\n",
        "        self.policy_net = PolicyNetwork(env.observation_space.shape[0], env.action_space.shape[0]).to(self.device)\n",
        "        \n",
        "        # copy params to target param\n",
        "        for target_param, param in zip(self.target_value_net.parameters(), self.value_net.parameters()):\n",
        "            target_param.data.copy_(param)\n",
        "            \n",
        "        # initialize optimizers \n",
        "        self.value_optimizer = optim.Adam(self.value_net.parameters(), lr=v_lr)\n",
        "        self.q1_optimizer = optim.Adam(self.q_net1.parameters(), lr=q_lr)\n",
        "        self.q2_optimizer = optim.Adam(self.q_net2.parameters(), lr=q_lr)\n",
        "        self.policy_optimizer = optim.Adam(self.policy_net.parameters(), lr=policy_lr)\n",
        "\n",
        "        self.replay_buffer = BasicBuffer(buffer_maxlen)\n",
        "        \n",
        "    def get_action(self, state):\n",
        "        state = torch.FloatTensor(state).unsqueeze(0).to(self.device)\n",
        "        mean, log_std = self.policy_net.forward(state)\n",
        "        std = log_std.exp()\n",
        "        \n",
        "        normal = Normal(mean, std)\n",
        "        z = normal.sample()\n",
        "        action = torch.tanh(z)\n",
        "        action = action.cpu().detach().squeeze(0).numpy()\n",
        "        \n",
        "        return self.rescale_action(action)\n",
        "    \n",
        "    def rescale_action(self, action):\n",
        "        return action * (self.action_range[1] - self.action_range[0]) / 2.0 +\\\n",
        "            (self.action_range[1] + self.action_range[0]) / 2.0\n",
        "   \n",
        "    def update(self, batch_size):\n",
        "        states, actions, rewards, next_states, dones = self.replay_buffer.sample(batch_size)\n",
        "        states = torch.FloatTensor(states).to(self.device)\n",
        "        actions = torch.FloatTensor(actions).to(self.device)\n",
        "        rewards = torch.FloatTensor(rewards).to(self.device)\n",
        "        next_states = torch.FloatTensor(next_states).to(self.device)\n",
        "        dones = torch.FloatTensor(dones).to(self.device)\n",
        "        dones = dones.view(dones.size(0), -1)\n",
        "        \n",
        "        next_actions, next_log_pi = self.policy_net.sample(next_states)\n",
        "        next_q1 = self.q_net1(next_states, next_actions)\n",
        "        next_q2 = self.q_net2(next_states, next_actions)\n",
        "        next_v = self.target_value_net(next_states)\n",
        "        \n",
        "        # value Loss\n",
        "        next_v_target = torch.min(next_q1, next_q2) - next_log_pi\n",
        "        curr_v = self.value_net.forward(states)\n",
        "        v_loss = F.mse_loss(curr_v, next_v_target.detach())\n",
        "        \n",
        "        # q loss\n",
        "        curr_q1 = self.q_net1.forward(states, actions)\n",
        "        curr_q2 = self.q_net2.forward(states, actions)\n",
        "        expected_q = rewards + (1 - dones) * self.gamma * next_v\n",
        "        q1_loss = F.mse_loss(curr_q1, expected_q.detach())\n",
        "        q2_loss = F.mse_loss(curr_q2, expected_q.detach())\n",
        "        \n",
        "        # update value network and q networks\n",
        "        self.value_optimizer.zero_grad()\n",
        "        v_loss.backward()\n",
        "        self.value_optimizer.step()\n",
        "        \n",
        "        self.q1_optimizer.zero_grad()\n",
        "        q1_loss.backward()\n",
        "        self.q1_optimizer.step()\n",
        "        \n",
        "        self.q2_optimizer.zero_grad()\n",
        "        q2_loss.backward()\n",
        "        self.q2_optimizer.step()\n",
        "        \n",
        "        #delayed update for policy net and target value nets\n",
        "        if self.update_step % self.delay_step == 0:\n",
        "            new_actions, log_pi = self.policy_net.sample(states)\n",
        "            min_q = torch.min(\n",
        "                self.q_net1.forward(states, new_actions),\n",
        "                self.q_net2.forward(states, new_actions)\n",
        "            )\n",
        "            policy_loss = (log_pi - min_q).mean()\n",
        "            \n",
        "            self.policy_optimizer.zero_grad()\n",
        "            policy_loss.backward()\n",
        "            self.policy_optimizer.step()\n",
        "        \n",
        "            # target networks\n",
        "            for target_param, param in zip(self.target_value_net.parameters(), self.value_net.parameters()):\n",
        "                target_param.data.copy_(self.tau * param + (1 - self.tau) * target_param)\n",
        "    \n",
        "        self.update_step += 1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1wWTGaoGSbWS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "env = gym.make(\"Pendulum-v0\")\n",
        "\n",
        "# SAC 2018 Params\n",
        "tau = 0.005\n",
        "gamma = 0.99\n",
        "value_lr = 3e-3\n",
        "q_lr = 3e-3\n",
        "policy_lr = 3e-3\n",
        "buffer_maxlen = 1000000\n",
        "\n",
        "# 2018 agent\n",
        "agent = SACAgent(env, gamma, tau, value_lr, q_lr, policy_lr, buffer_maxlen)\n",
        "\n",
        "# train\n",
        "episode_rewards = mini_batch_train(env, agent, 50, 500, 64)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}